///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file agp.h
///
/// \brief   header file for functions related to the antisymmetrized geminal power wavefunction
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef PCPS_AGP_HEADER
#define PCPS_AGP_HEADER

#include <src/pcps.h>
#include <src/input.h>
#include <src/wavefunction.h>
#include <src/blas_lapack.h>
#include <src/mantissa_exponent.h>

namespace pcps {

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   namespace for functions related to the antisymmetrized geminal power wavefunction
///
///////////////////////////////////////////////////////////////////////////////////////////////////

namespace agp {

  template <class SCALAR> bool check_sizes(const pcps::Wavefunction<SCALAR> & wfn,
                                           const int nap,
                                           const int nas,
                                           const bool use_pm);

  void populate_site_lists(const int nap,
                           const int nas,
                           const int * const config,
                           int * const a_sites,
                           int * const b_sites);

  template <class SCALAR> void build_occupied_matrix(const int nap,
                                                     const int nas,
                                                     const SCALAR * const gamma,
                                                     const int * const a_sites,
                                                     const int * const b_sites,
                                                     SCALAR * const mat);

  template <class SCALAR> pcps::MantExp<SCALAR> get_det_value(const pcps::Wavefunction<SCALAR> & wfn,
                                                              const int nap,
                                                              const int nas,
                                                              const int * const config,
                                                              const SCALAR * const gamma,
                                                              int * const a_sites,
                                                              int * const b_sites,
                                                              SCALAR * const work_mat,
                                                              int * const ipiv,
                                                              const bool use_pm);

  //-------------------------------------------------------------------------------
  // pcps::agp::init_matrices -- Initializes the gamma matrix and the inverse of
  //                             the agp matrix.
  //                             Returns the value of the agp matrix's determinant.
  //-------------------------------------------------------------------------------

                                                              // exit status  array size    description
  template <class SCALAR>
  void init_matrices(const pcps::Wavefunction<SCALAR> & wfn,  // unchanged                  the wavefunction object
                     const int nap,                           // unchanged                  number of alpha particles
                     const int nas,                           // unchanged                  number of alpha lattice sites
                     const int * const config,                // unchanged      2*nas       lattice configuration
                     SCALAR & det_value,                      // overwritten                on exit, the value of the agp matrix's determinant
                     int * const a_sites,                     // overwritten    nap         on exit, the occupied alpha sites (numbered 1 to nas)
                     int * const b_sites,                     // overwritten    nap         on exit, the occupied beta sites (numbered 1 to nas)
                     SCALAR * const gamma,                    // overwritten    nas*nas     on exit, the gamma matrix
                     SCALAR * const inv_mat,                  // overwritten    nap*nap     on exit, the inverse of the agp matrix
                     SCALAR * const agp_mat,                  // overwritten    nap*nap     work space for computing the agp matrix
                     int * const ipiv,                        // overwritten    nap         work space for pivoting during the LU decomposition
                     int * const inv_perm,                    // overwritten    nap         work space for inverting the LU pivoting
                     const bool use_pm);                      // unchanged                  whether to use the wavefunction's pairing matrix

  //-------------------------------------------------------------------------------
  // pcps::agp::det_ratio -- Computes the ratio of the AGP determinants for two
  //                         configurations that differ at the specified sites.
  //                         Returns   det(new config) / det(old config)
  //-------------------------------------------------------------------------------

                                                          // exit status  array size    description
  template <class SCALAR>
  SCALAR det_ratio(const int nap,                         // unchanged                  number of alpha particles
                   const int nas,                         // unchanged                  number of alpha lattice sites
                   const int nds,                         // unchanged                  the number of sites that differ between the configurations
                   const int * const diff_sites,          // unchanged      nds         the indices of the differing sites ( 0 <= index < 2*nas )
                   const int * const new_config,          // unchanged      2*nas       new lattice configuration
                   const int * const old_config,          // unchanged      2*nas       old lattice configuration
                   const SCALAR * const gamma,            // unchanged      nas*nas     the gamma matrix
                   const SCALAR * const inv_mat,          // unchanged      nap*nap     the inverse of the agp matrix
                   const int * const a_sites,             // unchanged      nap         the occupied alpha sites ( numbered 0 to nas-1 )
                   const int * const b_sites,             // unchanged      nap         the occupied beta sites  ( numbered 0 to nas-1 )
                   SCALAR * const wv1,                    // overwritten    nap         work space
                   SCALAR * const wv2,                    // overwritten    nap         work space
                   SCALAR * const wv3);                   // overwritten    nap         work space

  //-------------------------------------------------------------------------------
  // pcps::agp::update -- Updates the inverse matrix and alpha/beta site lists for
  //                      a new lattice configuration, which differs at the
  //                      specified sites from the old configuration.
  //-------------------------------------------------------------------------------

                                                          // exit status  array size    description
  template <class SCALAR>
  void update(const int nap,                              // unchanged                  number of alpha particles
              const int nas,                              // unchanged                  number of alpha lattice sites
              const int nds,                              // unchanged                  the number of sites that differ between the configurations
              const int * const diff_sites,               // unchanged      nds         the indices of the differing sites ( 0 <= index < 2*nas )
              const int * const new_config,               // unchanged      2*nas       new lattice configuration
              const int * const old_config,               // unchanged      2*nas       old lattice configuration
              const SCALAR * const gamma,                 // unchanged      nas*nas     the gamma matrix
              SCALAR * const inv_mat,                     // overwritten    nap*nap     the inverse of the agp matrix
              int * const a_sites,                        // overwritten    nap         the occupied alpha sites ( numbered 0 to nas-1 )
              int * const b_sites,                        // overwritten    nap         the occupied beta sites  ( numbered 0 to nas-1 )
              SCALAR * const wv1,                         // overwritten    nap         work space
              SCALAR * const wv2,                         // overwritten    nap         work space
              SCALAR * const wv3);                        // overwritten    nap         work space

  //-------------------------------------------------------------------------------
  // pcps::agp::der_ratios -- Computes the ratios between the AGP determinant's
  //                          derivatives and the AGP determinant's value for
  //                          the requested variables.
  //-------------------------------------------------------------------------------

                                                          // exit status  array size    description
  template <class SCALAR>
  void der_ratios(const pcps::Wavefunction<SCALAR> & wfn, // unchanged                  the wavefunction object
                  const int nap,                          // unchanged                  number of alpha particles
                  const int nas,                          // unchanged                  number of alpha lattice sites
                  const bool do_agp,                      // unchanged                  whether to compute derivatives for agp weights
                  const bool do_orb,                      // unchanged                  whether to compute derivatives for orbital coefficients
                  const bool do_par,                      // unchanged                  whether to compute derivatives for pairing matrix elements
                  const SCALAR * const gamma,             // unchanged      nas*nas     the gamma matrix
                  const SCALAR * const inv_mat,           // unchanged      nap*nap     the inverse of the agp matrix
                  const int * const a_sites,              // unchanged      nap         the occupied alpha sites ( numbered 0 to nas-1 )
                  const int * const b_sites,              // unchanged      nap         the occupied beta sites  ( numbered 0 to nas-1 )
                  SCALAR * const agp_der,                 // overwritten    nas         derivative ratios for agp weights
                  SCALAR * const orb_der,                 // overwritten    nas*nas     derivative ratios for orbital coefficients
                  SCALAR * const par_der);                // overwritten    nas*nas     derivative ratios for pairing matrix elements

  //-------------------------------------------------------------------------------
  // pcps::agp::row_diff -- Computes the difference between two rows of the
  //                        AGP matrix.
  //-------------------------------------------------------------------------------

                                                              // exit status  array size    description
  template <class SCALAR>
  inline void row_diff(const int nap,                         // unchanged                  number of alpha particles
                       const int nas,                         // unchanged                  number of alpha lattice sites
                       const int new_row,                     // unchanged                  the index of the new row ( 0 <= new_row < nas )
                       const int old_row,                     // unchanged                  the index of the old row ( 0 <= old_row < nas )
                       const SCALAR * const gamma,            // unchanged      nas*nas     the gamma matrix
                       const int * const b_sites,             // unchanged      nap         the occupied beta sites ( numbered 0 to nas-1 )
                       SCALAR * const diff)                   // overwritten    nap         on exit, the difference between the rows
  {
    for (int i = 0; i < nap; i++)
      diff[i] = gamma[ new_row * nas + b_sites[i] ] - gamma[ old_row * nas + b_sites[i] ];
  }

  //-------------------------------------------------------------------------------
  // pcps::agp::col_diff -- Computes the difference between two columns of the
  //                        AGP matrix.
  //-------------------------------------------------------------------------------

                                                              // exit status  array size    description
  template <class SCALAR>
  inline void col_diff(const int nap,                         // unchanged                  number of alpha particles
                       const int nas,                         // unchanged                  number of alpha lattice sites
                       const int new_col,                     // unchanged                  the index of the new col ( 0 <= new_col < nas )
                       const int old_col,                     // unchanged                  the index of the old col ( 0 <= old_col < nas )
                       const SCALAR * const gamma,            // unchanged      nas*nas     the gamma matrix
                       const int * const a_sites,             // unchanged      nap         the occupied alpha sites ( numbered 0 to nas-1 )
                       SCALAR * const diff)                   // overwritten    nap         on exit, the difference between the columns
  {
    for (int i = 0; i < nap; i++)
      diff[i] = gamma[ a_sites[i] * nas + new_col ] - gamma[ a_sites[i] * nas + old_col ];
  }

  //-------------------------------------------------------------------------------
  // pcps::agp::all_different -- Check that the elements of an array are distinct.
  //-------------------------------------------------------------------------------

                                                              // exit status  array size    description

  inline bool all_different(const int length,                 // unchanged                  length of the array
                            const int * const a)              // unchanged    length        the array to check
  {
    for (int i = 0; i < length; i++)
      for (int j = i+1; j < length; j++)
        if (a[i] == a[j])
          return false;
    return true;
  }

  //-------------------------------------------------------------------------------
  // pcps::agp::find_index -- Find the row or column index of the AGP matrix
  //                          corresponding to an occupied site.
  //-------------------------------------------------------------------------------

                                                              // exit status  array size    description

  inline int find_index(const int nap,                        // unchanged                  number of alpha particles
                        const int site,                       // unchanged                  the site whose index is to be found ( 0 <= sites < nas )
                        const int * const occ_sites)          // unchanged    nap           the list of occupied alpha or beta sites
  {
    assert( pcps::agp::all_different(nap, occ_sites) );
    int retval = -1;
    for (int i = 0; i < nap; i++) {
      if (occ_sites[i] == site) {
        retval = i;
        break;
      }
    }
    assert( retval >= 0 && retval < nap );
    return retval;
  }

  //-------------------------------------------------------------------------------
  // pcps::agp::get_diff_occs -- For two lattice configurations that differ at the
  //                             specified sites, this function determines which
  //                             sites are newly occupied and which are newly empty
  //-------------------------------------------------------------------------------

                                                              // exit status  array size    description

  inline void get_diff_occs(const int nds,                    // unchanged                  the number of sites that differ between the configurations
                            const int * const diff_sites,     // unchanged    nds           the indices of the differing sites ( 0 <= index < 2*nas )
                            const int * const new_config,     // unchanged    2*nas         the new lattice configuration
                            const int * const old_config,     // unchanged    2*nas         the old lattice configuration
                            int & n_a_diff,                   // overwritten                on exit, the number of differing alpha sites
                            int & n_b_diff,                   // overwritten                on exit, the number of differing beta sites
                            int * const a_new_occ_sites,      // overwritten  2             on exit, the newly occupied alpha sites
                            int * const b_new_occ_sites,      // overwritten  2             on exit, the newly occupied beta sites
                            int * const a_old_occ_sites,      // overwritten  2             on exit, the newly empty alpha sites
                            int * const b_old_occ_sites)      // overwritten  2             on exit, the newly empty beta sites
  {

    // declare some indices
    int i, j;

    // initialize the counters
    n_a_diff = 0;
    n_b_diff = 0;

    // newly occupied alpha sites
    for (i = j = 0; i < nds; i++)
      if ( diff_sites[i] % 2 == 0 && new_config[diff_sites[i]] == 1 )
        a_new_occ_sites[j++] = diff_sites[i] / 2;
    assert( j <= nds / 2 );
    n_a_diff = j;

    // previously occupied alpha sites
    for (i = j = 0; i < nds; i++)
      if ( diff_sites[i] % 2 == 0 && old_config[diff_sites[i]] == 1 )
        a_old_occ_sites[j++] = diff_sites[i] / 2;
    assert( j == n_a_diff );
    n_a_diff += j;

    // newly occupied beta sites
    for (i = j = 0; i < nds; i++)
      if ( diff_sites[i] % 2 == 1 && new_config[diff_sites[i]] == 1 )
        b_new_occ_sites[j++] = diff_sites[i] / 2;
    assert( j <= nds / 2 );
    n_b_diff = j;

    // previously occupied beta sites
    for (i = j = 0; i < nds; i++)
      if ( diff_sites[i] % 2 == 1 && old_config[diff_sites[i]] == 1 )
        b_old_occ_sites[j++] = diff_sites[i] / 2;
    assert( j == n_b_diff );
    n_b_diff += j;

    // check that the numbers of different alpha and beta sites are acceptable
    assert( nds != 0 || (n_a_diff == 0 && n_b_diff == 0) );
    assert( nds != 2 || (n_a_diff == 2 && n_b_diff == 0) || (n_a_diff == 0 && n_b_diff == 2) );
    assert( nds != 4 || (n_a_diff == 2 && n_b_diff == 2) || (n_a_diff == 4 && n_b_diff == 0) || (n_a_diff == 0 && n_b_diff == 4) );

    // check that the sites' occupations are correct
    assert( n_a_diff < 2 || ( 1 == new_config[2*a_new_occ_sites[0]+0] && 0 == old_config[2*a_new_occ_sites[0]+0] ) );
    assert( n_a_diff < 2 || ( 1 == old_config[2*a_old_occ_sites[0]+0] && 0 == new_config[2*a_old_occ_sites[0]+0] ) );
    assert( n_b_diff < 2 || ( 1 == new_config[2*b_new_occ_sites[0]+1] && 0 == old_config[2*b_new_occ_sites[0]+1] ) );
    assert( n_b_diff < 2 || ( 1 == old_config[2*b_old_occ_sites[0]+1] && 0 == new_config[2*b_old_occ_sites[0]+1] ) );
    assert( n_a_diff < 4 || ( 1 == new_config[2*a_new_occ_sites[1]+0] && 0 == old_config[2*a_new_occ_sites[1]+0] ) );
    assert( n_a_diff < 4 || ( 1 == old_config[2*a_old_occ_sites[1]+0] && 0 == new_config[2*a_old_occ_sites[1]+0] ) );
    assert( n_b_diff < 4 || ( 1 == new_config[2*b_new_occ_sites[1]+1] && 0 == old_config[2*b_new_occ_sites[1]+1] ) );
    assert( n_b_diff < 4 || ( 1 == old_config[2*b_old_occ_sites[1]+1] && 0 == new_config[2*b_old_occ_sites[1]+1] ) );

  }

  //-------------------------------------------------------------------------------
  // pcps::agp::row_update -- Updates the inverse matrix and alpha site list after
  //                          a row replacement in the agp matrix.
  //-------------------------------------------------------------------------------

                                                              // exit status  array size    description
  template <class SCALAR>
  inline void row_update(const int nap,                       // unchanged                  number of alpha particles
                         const int nas,                       // unchanged                  number of alpha lattice sites
                         const int new_occ_site,              // unchanged                  the newly occupied site ( 0 <= new_occ_site < nas-1 )
                         const int old_occ_site,              // unchanged                  the newly empty site    ( 0 <= old_occ_site < nas-1 )
                         const SCALAR * const gamma,          // unchanged      nas*nas     the gamma matrix
                         const int * const b_sites,           // unchanged      nap         the occupied beta sites  ( numbered 0 to nas-1 )
                         int * const a_sites,                 // overwritten    nap         the occupied alpha sites ( numbered 0 to nas-1 )
                         SCALAR * const inv_mat,              // overwritten    nap*nap     the inverse of the agp matrix
                         SCALAR * const wv1,                  // overwritten    nap         work space
                         SCALAR * const wv2,                  // overwritten    nap         work space
                         SCALAR * const wv3)                  // overwritten    nap         work space
  {

    // check that the new and old sites are different
    assert( new_occ_site != old_occ_site );

    // get the index of the row to update
    const int row = pcps::agp::find_index(nap, old_occ_site, a_sites);

    // get the difference between the new and old row
    pcps::agp::row_diff(nap, nas, new_occ_site, old_occ_site, gamma, b_sites, wv1);

    // get the necessary column of the inverse matrix
    for (int i = 0; i < nap; i++)
      wv2[i] = inv_mat[ nap * i + row ];

    // compute the determinant ratio
    const SCALAR ratio = pcps::unity<SCALAR>() + pcps::xdot(nap, wv1, 1, wv2, 1);

    // compute the action of the inverse matrix on the row difference
    pcps::xgemm('N', 'N', nap, 1, nap, pcps::unity<SCALAR>(), inv_mat, nap, wv1, nap, pcps::zero<SCALAR>(), wv3, nap);

    // update the inverse matrix
    pcps::xgemm('N', 'N', nap, nap, 1, -1.0 / ratio, wv3, nap, wv2, 1, pcps::unity<SCALAR>(), inv_mat, nap);

    // update the occupied site list
    a_sites[row] = new_occ_site;

  }

  //-------------------------------------------------------------------------------
  // pcps::agp::col_update -- Updates the inverse matrix and beta site list after
  //                          a column replacement in the agp matrix.
  //-------------------------------------------------------------------------------

                                                              // exit status  array size    description
  template <class SCALAR>
  inline void col_update(const int nap,                       // unchanged                  number of alpha particles
                         const int nas,                       // unchanged                  number of alpha lattice sites
                         const int new_occ_site,              // unchanged                  the newly occupied site ( 0 <= new_occ_site < nas-1 )
                         const int old_occ_site,              // unchanged                  the newly empty site    ( 0 <= old_occ_site < nas-1 )
                         const SCALAR * const gamma,          // unchanged      nas*nas     the gamma matrix
                         const int * const a_sites,           // unchanged      nap         the occupied alpha sites ( numbered 0 to nas-1 )
                         int * const b_sites,                 // overwritten    nap         the occupied beta sites  ( numbered 0 to nas-1 )
                         SCALAR * const inv_mat,              // overwritten    nap*nap     the inverse of the agp matrix
                         SCALAR * const wv1,                  // overwritten    nap         work space
                         SCALAR * const wv2,                  // overwritten    nap         work space
                         SCALAR * const wv3)                  // overwritten    nap         work space
  {

    // check that the new and old sites are different
    assert( new_occ_site != old_occ_site );

    // get the index of the column to update
    const int col = pcps::agp::find_index(nap, old_occ_site, b_sites);

    // get the difference between the new and old column
    pcps::agp::col_diff(nap, nas, new_occ_site, old_occ_site, gamma, a_sites, wv1);

    // get the necessary row of the inverse matrix
    for (int i = 0; i < nap; i++)
      wv2[i] = inv_mat[ nap * col + i ];

    // compute the determinant ratio
    const SCALAR ratio = pcps::unity<SCALAR>() + pcps::xdot(nap, wv1, 1, wv2, 1);

    // compute the action of the inverse matrix on the column difference
    pcps::xgemm('T', 'N', nap, 1, nap, pcps::unity<SCALAR>(), inv_mat, nap, wv1, nap, pcps::zero<SCALAR>(), wv3, nap);

    // update the inverse matrix
    pcps::xgemm('N', 'N', nap, nap, 1, -1.0 / ratio, wv2, nap, wv3, 1, pcps::unity<SCALAR>(), inv_mat, nap);

    // update the occupied site list
    b_sites[col] = new_occ_site;

  }

} // end namespace agp
} // end namespace pcps

#endif
